
#include "km_match.h"
#include "utils.hpp"
namespace perception
{
    namespace algorithm
    {
        MatchKM::MatchKM(rclcpp::Node *private_nh) : private_nh_(private_nh) {}
        MatchKM::~MatchKM() {}
        void MatchKM::Process(std::vector<LidarObject> &input_objs,
                              std::vector<Match> &now_objs, float &euclidean_distance,
                              std::vector<bool> &g_matched,
                              std::vector<bool> &c_matched,
                              std::vector<std::pair<int, int>> &pairs)
        {
            std::vector<std::vector<int>> cost_matrix;
            // 得到匹配的距离及id
            for (int g = 0; g < input_objs.size(); ++g)
            {
                std::vector<int> tmp;
                for (int c = 0; c < now_objs.size(); ++c)
                {
                    // 增加一个预测 这个是随便弄的
                    // int dis = int((std::fabs(input_objs[g].corner_point.x -
                    // now_objs[c].cp.x) +
                    //                std::fabs(input_objs[g].corner_point.y -
                    //                now_objs[c].cp.y)) *
                    //               100);
                    int dis = int((std::fabs(input_objs[g].x - now_objs[c].x) +
                                   std::fabs(input_objs[g].y - now_objs[c].y)) *
                                  100);

                    tmp.emplace_back(dis);
                }
                cost_matrix.emplace_back(tmp);
            }

            auto mask_matrix = munkresAssignment(cost_matrix);
            for (int i = 0; i < mask_matrix.size(); i++)
            {
                for (int j = 0; j < mask_matrix[i].size(); j++)
                {
                    if (mask_matrix[i][j] == 1)
                    {
                        if (cost_matrix[i][j] < euclidean_distance * 100)
                        {
                            g_matched[i] = true;
                            c_matched[j] = true;
                            std::pair<int, int> p;
                            p.first = j;
                            p.second = i;
                            pairs.emplace_back(p);
                        }
                    }
                }
            }
        }
        void MatchKM::Process(std::vector<LidarObject> &input_objs,
                              std::vector<MatchCenter> &now_objs,
                              float &euclidean_distance, std::vector<bool> &g_matched,
                              std::vector<bool> &c_matched,
                              std::vector<std::pair<int, int>> &pairs)
        {
            std::vector<std::vector<int>> cost_matrix;
            // 得到匹配的距离及id
            // RCLCPP_INFO(private_nh_->get_logger(), "0000000");
            // RCLCPP_INFO(private_nh_->get_logger(), "gsize:%d,csize:%d",
            // input_objs.size(), now_objs.size());
            for (int g = 0; g < input_objs.size(); ++g)
            {
                std::vector<int> tmp;
                for (int c = 0; c < now_objs.size(); ++c)
                {
                    // 增加一个预测 这个是随便弄的
                    // int dis = int((std::fabs(input_objs[g].corner_point.x -
                    // now_objs[c].cp.x) +
                    //                std::fabs(input_objs[g].corner_point.y -
                    //                now_objs[c].cp.y)) *
                    //               100);
                    // RCLCPP_INFO(private_nh_->get_logger(), "gx:%f,gy:%f,cx:%f,cy:%f",
                    //             input_objs[g].x, input_objs[g].y, now_objs[c].x,
                    //             now_objs[c].y);
                    int dis = int((std::fabs(input_objs[g].x - now_objs[c].x) +
                                   std::fabs(input_objs[g].y - now_objs[c].y)) *
                                  100);
                    //   LidarObject obj;
                    // now_objs[c].GetObj(obj);
                    // RCLCPP_INFO(private_nh_->get_logger(), "111111");
                    // float iou;
                    // utils::CalcIou(input_objs[g], obj, iou);
                    // if (iou < 0.1) {
                    //   dis = int((euclidean_distance + 1) * 100);
                    // }
                    //   auto dx = now_objs[c].cvx * 0.1 + 2;
                    //   auto dy = now_objs[c].cvy * 0.1 + 2;
                    //   if (std::fabs(input_objs[g].x - now_objs[c].x) > dx ||
                    //       std::fabs(input_objs[g].y - now_objs[c].y) > dy)
                    //     dis = int((euclidean_distance + 1) * 100);
                    tmp.emplace_back(dis);
                }
                cost_matrix.emplace_back(tmp);
            }
            // RCLCPP_INFO(private_nh_->get_logger(), "111111");
            auto mask_matrix = munkresAssignment(cost_matrix);
            // RCLCPP_INFO(private_nh_->get_logger(), "222222");
            for (int i = 0; i < mask_matrix.size(); i++)
            {
                for (int j = 0; j < mask_matrix[i].size(); j++)
                {
                    if (mask_matrix[i][j] == 1)
                    {
                        if (cost_matrix[i][j] < euclidean_distance * 100)
                        {
                            g_matched[i] = true;
                            c_matched[j] = true;
                            std::pair<int, int> p;
                            p.first = j;
                            p.second = i;
                            pairs.emplace_back(p);
                        }
                    }
                }
            }
        }

        void MatchKM::step_one(int &step, std::vector<std::vector<int>> &cost_matrix,
                               std::vector<std::vector<int>> &mask_matrix)
        {
            int nrow = cost_matrix.size();
            int ncol = cost_matrix[0].size();
            int min_in_row;
            for (int i = 0; i < nrow; i++)
            {
                std::vector<int>::iterator min_in_row =
                    std::min_element(cost_matrix[i].begin(), cost_matrix[i].end());
                for (int j = 0; j < ncol; j++)
                {
                    cost_matrix[i][j] -= *min_in_row;
                }
            }
            step = 2;
        }

        void MatchKM::step_two(int &step, std::vector<std::vector<int>> &cost_matrix,
                               std::vector<std::vector<int>> &mask_matrix,
                               std::vector<int> &rowCover, std::vector<int> &colCover)
        {
            int nrow = cost_matrix.size();
            int ncol = cost_matrix[0].size();
            for (int i = 0; i < nrow; i++)
            {
                for (int j = 0; j < ncol; j++)
                {
                    if (cost_matrix[i][j] == 0 && rowCover[i] == 0 && colCover[j] == 0)
                    {
                        mask_matrix[i][j] = 1;
                        rowCover[i] = 1;
                        colCover[j] = 1;
                    }
                }
            }
            for (int i = 0; i < nrow; i++)
            {
                rowCover[i] = 0;
            }
            for (int j = 0; j < ncol; j++)
            {
                colCover[j] = 0;
            }
            step = 3;
        }

        void MatchKM::step_three(int &step, std::vector<std::vector<int>> &cost_matrix,
                                 std::vector<std::vector<int>> &mask_matrix,
                                 std::vector<int> &rowCover,
                                 std::vector<int> &colCover)
        {
            int nrow = cost_matrix.size();
            int ncol = cost_matrix[0].size();
            for (int i = 0; i < nrow; i++)
            {
                for (int j = 0; j < ncol; j++)
                {
                    if (mask_matrix[i][j] == 1)
                    {
                        colCover[j] = 1;
                    }
                }
            }

            int colcount = 0;
            for (int j = 0; j < ncol; j++)
            {
                if (colCover[j] == 1)
                {
                    colcount += 1;
                }
            }

            if (colcount >= ncol || colcount >= nrow)
            {
                step = 7;
            }
            else
            {
                step = 4;
            }
        }

        void MatchKM::find_a_zero(int &row, int &col, int &step,
                                  std::vector<std::vector<int>> &cost_matrix,
                                  std::vector<std::vector<int>> &mask_matrix,
                                  std::vector<int> &rowCover,
                                  std::vector<int> &colCover)
        {
            int i = 0;
            int j;
            bool done = false;
            row = -1;
            col = -1;

            int nrow = cost_matrix.size();
            int ncol = cost_matrix[0].size();
            while (!done)
            {
                j = 0;
                while (1)
                {
                    if (cost_matrix[i][j] == 0 && rowCover[i] == 0 && colCover[j] == 0)
                    {
                        row = i;
                        col = j;
                        done = true;
                    }
                    j++;
                    if (j >= ncol || done)
                    {
                        break;
                    }
                }
                i++;
                if (i >= nrow)
                {
                    done = true;
                }
            }
        }

        bool MatchKM::star_in_row(int &row, int &col,
                                  std::vector<std::vector<int>> &mask_matrix)
        {
            bool tmp = false;
            int ncol = mask_matrix[0].size();
            for (int j = 0; j < ncol; j++)
            {
                if (mask_matrix[row][j] == 1)
                {
                    tmp = true;
                }
            }
            return tmp;
        }

        void MatchKM::find_star_in_row(int &row, int &col,
                                       std::vector<std::vector<int>> &mask_matrix)
        {
            col = -1;
            int ncol = mask_matrix[0].size();
            for (int j = 0; j < ncol; j++)
            {
                if (mask_matrix[row][j] == 1)
                {
                    col = j;
                }
            }
        }

        void MatchKM::find_star_in_col(int &row, int &col,
                                       std::vector<std::vector<int>> &mask_matrix)
        {
            row = -1;
            int nrow = mask_matrix.size();
            for (int i = 0; i < nrow; i++)
            {
                if (mask_matrix[i][col] == 1)
                {
                    row = i;
                }
            }
        }

        void MatchKM::find_prime_in_row(int &row, int &col,
                                        std::vector<std::vector<int>> &mask_matrix)
        {
            int ncol = mask_matrix[0].size();
            for (int j = 0; j < ncol; j++)
            {
                if (mask_matrix[row][j] == 2)
                {
                    col = j;
                }
            }
        }

        void MatchKM::augment_path(int &path_count, std::vector<std::vector<int>> path,
                                   std::vector<std::vector<int>> &mask_matrix)
        {
            for (int p = 0; p < path_count; p++)
            {
                if (mask_matrix[path[p][0]][path[p][1]] == 1)
                {
                    mask_matrix[path[p][0]][path[p][1]] = 0;
                }
                else
                {
                    mask_matrix[path[p][0]][path[p][1]] = 1;
                }
            }
        }

        void MatchKM::clear_covers(std::vector<int> &rowCover,
                                   std::vector<int> &colCover)
        {
            for (int i = 0; i < rowCover.size(); i++)
            {
                rowCover[i] = 0;
            }
            for (int j = 0; j < colCover.size(); j++)
            {
                colCover[j] = 0;
            }
        }

        void MatchKM::erase_primes(std::vector<std::vector<int>> &mask_matrix)
        {
            int nrow = mask_matrix.size();
            int ncol = mask_matrix[0].size();
            for (int i = 0; i < nrow; i++)
            {
                for (int j = 0; j < ncol; j++)
                {
                    if (mask_matrix[i][j] == 2)
                    {
                        mask_matrix[i][j] = 0;
                    }
                }
            }
        }

        void MatchKM::step_four(int &step, std::vector<std::vector<int>> &cost_matrix,
                                std::vector<std::vector<int>> &mask_matrix,
                                std::vector<int> &rowCover, std::vector<int> &colCover,
                                std::vector<std::vector<int>> &path, int &path_row_0,
                                int &path_col_0)
        {
            int row = -1;
            int col = -1;
            bool done = false;

            while (!done)
            {
                find_a_zero(row, col, step, cost_matrix, mask_matrix, rowCover, colCover);
                if (row == -1)
                {
                    done = true;
                    step = 6;
                }
                else
                {
                    mask_matrix[row][col] = 2;
                    if (star_in_row(row, col, mask_matrix))
                    {
                        find_star_in_row(row, col, mask_matrix);
                        rowCover[row] = 1;
                        colCover[col] = 0;
                    }
                    else
                    {
                        done = true;
                        step = 5;
                        path_row_0 = row;
                        path_col_0 = col;
                    }
                }
            }
        }

        void MatchKM::step_five(int &step, std::vector<std::vector<int>> &cost_matrix,
                                std::vector<std::vector<int>> &mask_matrix,
                                std::vector<int> &rowCover, std::vector<int> &colCover,
                                std::vector<std::vector<int>> &path, int &path_row_0,
                                int &path_col_0)
        {
            bool done = false;
            int i = -1;
            int j = -1;

            int path_count = 1;
            path[path_count - 1][0] = path_row_0;
            path[path_count - 1][1] = path_col_0;

            while (!done)
            {
                find_star_in_col(i, path[path_count - 1][1], mask_matrix);

                if (i > -1)
                {

                    if (path_count >= path.size())
                    {
                        done = true;
                        break;
                    }
                    path_count += 1;
                    path[path_count - 1][0] = i;
                    path[path_count - 1][1] = path[path_count - 2][1];
                }
                else
                {

                    done = true;
                }

                if (!done)
                {
                    find_prime_in_row(path[path_count - 1][0], j, mask_matrix);

                    if (path_count >= path.size())
                    {
                        done = true;
                        break;
                    }

                    path_count += 1;

                    path[path_count - 1][0] = path[path_count - 2][0];
                    path[path_count - 1][1] = j;
                }
            }

            augment_path(path_count, path, mask_matrix);
            clear_covers(rowCover, colCover);
            erase_primes(mask_matrix);
            step = 3;
        }

        void MatchKM::find_smallest(int &minval,
                                    std::vector<std::vector<int>> &cost_matrix,
                                    std::vector<int> &rowCover,
                                    std::vector<int> &colCover)
        {
            for (int i = 0; i < rowCover.size(); i++)
            {
                for (int j = 0; j < colCover.size(); j++)
                {
                    if (rowCover[i] == 0 && colCover[j] == 0)
                    {
                        if (minval > cost_matrix[i][j])
                        {
                            minval = cost_matrix[i][j];
                        }
                    }
                }
            }
        }

        void MatchKM::step_six(int &step, std::vector<std::vector<int>> &cost_matrix,
                               std::vector<int> &rowCover, std::vector<int> &colCover)
        {
            int minval = 0x7fffffff;
            find_smallest(minval, cost_matrix, rowCover, colCover);

            for (int i = 0; i < rowCover.size(); i++)
            {
                for (int j = 0; j < colCover.size(); j++)
                {
                    if (rowCover[i] == 1)
                    {
                        cost_matrix[i][j] += minval;
                    }
                    if (colCover[j] == 0)
                    {
                        cost_matrix[i][j] -= minval;
                    }
                }
            }
            step = 4;
        }

        void MatchKM::step_seven()
        {
            // std::cout << "Munkres Assignment Algorithm finish!!!" << std::endl;
        }

        void MatchKM::process_step(bool &done, int &step,
                                   std::vector<std::vector<int>> cost_matrix_origin,
                                   std::vector<std::vector<int>> &cost_matrix,
                                   std::vector<std::vector<int>> &mask_matrix,
                                   std::vector<int> &rowCover,
                                   std::vector<int> &colCover,
                                   std::vector<std::vector<int>> &path, int &path_row_0,
                                   int &path_col_0)
        {
            // std::cout << "excute step: " << step << std::endl;
            switch (step)
            {
            case 1:
                step_one(step, cost_matrix, mask_matrix);
                break;
            case 2:
                step_two(step, cost_matrix, mask_matrix, rowCover, colCover);
                break;
            case 3:
                step_three(step, cost_matrix, mask_matrix, rowCover, colCover);
                break;
            case 4:
                step_four(step, cost_matrix, mask_matrix, rowCover, colCover, path,
                          path_row_0, path_col_0);
                break;
            case 5:
                step_five(step, cost_matrix, mask_matrix, rowCover, colCover, path,
                          path_row_0, path_col_0);
                break;
            case 6:
                step_six(step, cost_matrix, rowCover, colCover);
                break;
            case 7:
                step_seven();
                done = true;

                // int mini_cost_value = 0;
                // for (int i = 0; i < mask_matrix.size(); i++)
                // {
                //     for (int j = 0; j < mask_matrix[i].size(); j++)
                //     {
                //         if (mask_matrix[i][j] == 1)
                //         {
                //             std::cout << "use data on (" << i << ", " << j << ") value: "
                //             << cost_matrix_origin[i][j] << std::endl; mini_cost_value +=
                //             cost_matrix_origin[i][j];
                //         }
                //     }
                // }
                // std::cout << "mini cost:" << mini_cost_value << std::endl;
                break;
            }
        }

        void MatchKM::showCostMatrix(std::vector<std::vector<int>> &cost_matrix)
        {
            std::cout << "__________cost_________" << std::endl;
            for (int i = 0; i < cost_matrix.size(); i++)
            {
                for (int j = 0; j < cost_matrix[i].size(); j++)
                {
                    std::cout << cost_matrix[i][j] << " ";
                }
                std::cout << std::endl;
            }
        }

        void MatchKM::showMaskMatrix(std::vector<std::vector<int>> &mask_matrix)
        {
            std::cout << "__________mask_________" << std::endl;
            for (int i = 0; i < mask_matrix.size(); i++)
            {
                for (int j = 0; j < mask_matrix[i].size(); j++)
                {
                    std::cout << mask_matrix[i][j] << " ";
                }
                std::cout << std::endl;
            }
        }

        std::vector<std::vector<int>>
        MatchKM::munkresAssignment(std::vector<std::vector<int>> &cost_matrix)
        {

            std::vector<std::vector<int>> cost_matrix_origin = cost_matrix;

            std::vector<std::vector<int>> mask_matrix;
            for (int i = 0; i < cost_matrix.size(); i++)
            {
                std::vector<int> a;
                for (int j = 0; j < cost_matrix[i].size(); j++)
                {
                    a.push_back(0);
                }
                mask_matrix.push_back(a);
            }

            std::vector<int> rowCover;
            for (int i = 0; i < cost_matrix.size(); i++)
            {
                rowCover.push_back(0);
            }
            std::vector<int> colCover;
            for (int i = 0; i < cost_matrix[0].size(); i++)
            {
                colCover.push_back(0);
            }

            std::vector<std::vector<int>> path;
            for (int i = 0; i < cost_matrix.size(); i++)
            {
                std::vector<int> a;
                a.push_back(0);
                a.push_back(0);
                path.push_back(a);
            }

            bool done = false;
            int step = 1;
            int path_row_0, path_col_0;
            while (!done)
            {
                // showCostMatrix(cost_matrix);
                // showMaskMatrix(mask_matrix);

                process_step(done, step, cost_matrix_origin, cost_matrix, mask_matrix,
                             rowCover, colCover, path, path_row_0, path_col_0);
            }
            return mask_matrix;
        }
    } // namespace algorithm
} // namespace perception
